# coding=UTF-8
import tensorflow as tf
import os.path
import argparse
 
CHECKPOINT_PATH = "mobilenet_v1_1.0_224/mobilenet_v1_1.0_224.ckpt"
CHECKPOINT_PATH2 = "models/model.ckpt"
MODEL_NAME = "models/frozen_model.pb"
 
def freeze_graph():
    if tf.gfile.IsDirectory(CHECKPOINT_PATH):
      checkpoint_path = tf.train.latest_checkpoint(CHECKPOINT_PATH)
    else:
      checkpoint_path = CHECKPOINT_PATH
    with tf.Session() as sess:
        saver = tf.train.import_meta_graph(checkpoint_path + '.meta')
        saver.restore(sess, CHECKPOINT_PATH2) #恢复图并得到数据
        tensor = sess.graph.get_tensor_by_name('input:0')
        print(tensor)
        output_graph_def = tf.graph_util.convert_variables_to_constants(  #模型持久化，将变量值固定
            sess,
            input_graph_def=sess.graph.as_graph_def(),
            output_node_names=[
                'MobilenetV1/Predictions/Reshape_1']
        )
        with tf.gfile.GFile(MODEL_NAME, "wb") as f: #保存模型
            f.write(output_graph_def.SerializeToString()) #序列化输出

if __name__ == '__main__':
    # parser = argparse.ArgumentParser()
    # parser.add_argument("model_folder", type=str, help="input ckpt model dir") #命令行解析，help是提示符，type是输入的类型，
    # # 这里运行程序时需要带上模型ckpt的路径，不然会报 error: too few arguments
    # aggs = parser.parse_args()
    # freeze_graph(aggs.model_folder)
    freeze_graph()
 